#include "PropertyMapper.h"
#include "FileIO.h"
#include <fstream>
#include <iomanip>
#include <algorithm>

size_t PropertyMapper::K_NEAREST_POINTS = 8;  // 默认使用8个最近邻点

PropertyMapper::PropertyMapper() {}

void PropertyMapper::process(
    const std::string& gridFile,
    const std::string& propsFile,
    const std::vector<Point3D>& targetNodes,
    const std::vector<Tetrahedron>& targetElements,
    std::vector<CellProperty>& mappedProps) {
    
    // 1. 读取源网格数据
    GridSize size;
    std::vector<std::vector<double>> coord;
    std::vector<double> zcorn;
    FileIO::readGridFile(gridFile, size, coord, zcorn);
    
    // 2. 计算源网格中心点
    auto sourceCenters = calculateCornerGridCenters(size, coord, zcorn);
    
    // 3. 读取源网格属性
    GridProperties props;
    FileIO::readPropertiesFile(propsFile, props);
    
    // 4. 组织源网格属性数据
    std::vector<CellProperty> sourceProps(sourceCenters.size());
    for (size_t i = 0; i < sourceProps.size(); ++i) {
        sourceProps[i].youngModulos = props.youngmod[i];
        sourceProps[i].poissonRatio = props.poissonr[i];
        sourceProps[i].rockDensity = props.rockden[i];
        sourceProps[i].porosity = props.poro[i];
    }
    
    // 5. 计算目标网格中心点
    std::vector<Point3D> targetCenters;
    targetCenters.reserve(targetElements.size());
    for (const auto& element : targetElements) {
        targetCenters.push_back(calculateTetCenter(targetNodes, element));
    }
    
    // 6. 执行插值
    performInterpolation(sourceCenters, sourceProps, targetCenters, mappedProps);
}

Point3D PropertyMapper::calculateTetCenter(
    const std::vector<Point3D>& nodes,
    const Tetrahedron& tet) const {
    const auto& p1 = nodes[tet.v1];
    const auto& p2 = nodes[tet.v2];
    const auto& p3 = nodes[tet.v3];
    const auto& p4 = nodes[tet.v4];
    
    return Point3D(
        (p1.x + p2.x + p3.x + p4.x) / 4.0,
        (p1.y + p2.y + p3.y + p4.y) / 4.0,
        (p1.z + p2.z + p3.z + p4.z) / 4.0
    );
}

std::vector<Point3D> PropertyMapper::calculateCornerGridCenters(
    const GridSize& size,
    const std::vector<std::vector<double>>& coord,
    const std::vector<double>& zcorn) {
    
    std::vector<Point3D> centers;
    centers.reserve(size.Nx * size.Ny * size.Nz);
    
    // 重组ZCORN数据为3D数组
    std::vector<std::vector<std::vector<double>>> zcorn_3d(
        2 * size.Nz,
        std::vector<std::vector<double>>(
            2 * size.Ny,
            std::vector<double>(2 * size.Nx)
        )
    );
    
    size_t idx = 0;
    for (int k = 0; k < 2 * size.Nz; ++k)
        for (int j = 0; j < 2 * size.Ny; ++j)
            for (int i = 0; i < 2 * size.Nx; ++i)
                zcorn_3d[k][j][i] = zcorn[idx++];
    
    // 计算每个单元的中心点
    #pragma omp parallel for collapse(3)
    for (int k = 0; k < size.Nz; ++k) {
        for (int j = 0; j < size.Ny; ++j) {
            for (int i = 0; i < size.Nx; ++i) {
                Point3D center(0, 0, 0);
                int count = 0;
                
                // 计算8个顶点的平均值
                for (int kk = 0; kk <= 1; ++kk)
                    for (int jj = 0; jj <= 1; ++jj)
                        for (int ii = 0; ii <= 1; ++ii) {
                            center.x += coord[(j+jj)*(size.Nx+1) + (i+ii)][0];
                            center.y += coord[(j+jj)*(size.Nx+1) + (i+ii)][1];
                            center.z += zcorn_3d[2*k+kk][2*j+jj][2*i+ii];
                            count++;
                        }
                
                center.x /= count;
                center.y /= count;
                center.z /= count;
                
                #pragma omp critical
                centers.push_back(center);
            }
        }
    }
    
    return centers;
}

void PropertyMapper::performInterpolation(
    const std::vector<Point3D>& sourceCenters,
    const std::vector<CellProperty>& sourceProps,
    const std::vector<Point3D>& targetCenters,
    std::vector<CellProperty>& mappedProps) {
    
    // 构建KD树
    std::vector<std::pair<Point3D, size_t>> points;
    points.reserve(sourceCenters.size());
    for (size_t i = 0; i < sourceCenters.size(); ++i) {
        points.emplace_back(sourceCenters[i], i);
    }
    kdtree_.build(points);
    
    // 执行并行插值
    mappedProps.resize(targetCenters.size());
    #pragma omp parallel for schedule(dynamic)
    for (int i = 0; i < static_cast<int>(targetCenters.size()); ++i) {
        std::vector<std::pair<double, size_t>> neighbors;
        kdtree_.findKNearest(targetCenters[i], K_NEAREST_POINTS, neighbors);
        mappedProps[i] = calculateInterpolation(targetCenters[i], neighbors, sourceProps);
    }
}

CellProperty PropertyMapper::calculateInterpolation(
    const Point3D& target,
    const std::vector<std::pair<double, size_t>>& neighbors,
    const std::vector<CellProperty>& sourceProps) {
    
    double sum_weight = 0.0;
    CellProperty result = {0.0, 0.0, 0.0, 0.0};
    
    for (const auto& neighbor : neighbors) {
        double weight = 1.0 / (neighbor.first * neighbor.first);
        sum_weight += weight;
        
        const auto& prop = sourceProps[neighbor.second];
        result.youngModulos += weight * prop.youngModulos;
        result.poissonRatio += weight * prop.poissonRatio;
        result.rockDensity += weight * prop.rockDensity;
        result.porosity += weight * prop.porosity;
    }
    
    // 归一化
    result.youngModulos /= sum_weight;
    result.poissonRatio /= sum_weight;
    result.rockDensity /= sum_weight;
    result.porosity /= sum_weight;
    
    return result;
} 